/*
  Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

  This file is part of Sunrise.

  Sunrise is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 3 of the License, or
  (at your option) any later version.

  Sunrise is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.
  
*/

#include<iostream>
#include "CCfits/CCfits"
#include "counter.h"
#include "polychromatic_grid.h"
#include "optical.h"
#include "scatterer.h"
#include "constants.h"
#include "dust_model.h"
#include "emission.h"
#include "emission-fits.h"
#include "full_sed_emergence.h"
#include "full_sed_grid.h"
#include "emergence-fits.h"
#include "mcrx-units.h"
#include "grid_factory.h"
#include "grid-fits.h"
#include "shoot.h"
#include "fits-utilities.h"
#include "boost/shared_ptr.hpp"
#include "density_generator.h"
#include "preferences.h"
#include "model.h"
#include "wd01_Brent_PAH_grain_model.h"
#include "wd01_grain_model.h"
#include "xfer.h"
#include "ir_grid.h"
#include "equilibrium.h"

using namespace mcrx;
using namespace CCfits;
using namespace std; 

typedef local_random T_rng_policy;
typedef scatterer<polychromatic_scatterer_policy, T_rng_policy> T_scatterer;
typedef T_scatterer::T_lambda T_lambda;
typedef T_scatterer::T_biaser T_biaser;
typedef dust_model<T_scatterer, cumulative_sampling, T_rng_policy> T_dust_model;
typedef full_sed_emergence T_emergence;

typedef full_sed_grid<adaptive_grid> T_grid;
typedef ir_grid<adaptive_grid, cumulative_sampling, local_random> T_ir_grid;
typedef emission<polychromatic_policy, T_rng_policy> T_emission;

class uniform_factory: public grid_factory<T_grid::T_data> {
public:
  const T_densities rho_;
  const int n_lambda;

  uniform_factory (const T_densities& rho, int nl):
    rho_ (rho), n_lambda (nl) {};
  
  virtual bool refine_cell_p (const T_cell& c, int level) {
    return false;
  };
  virtual bool unify_cell_p (T_grid::const_local_iterator begin,
			     T_grid::const_local_iterator end) {
    return false;
  };
  virtual T_data get_data (const T_cell& c) {
    return T_data(rho_, vec3d(0,0,0), T_lambda(n_lambda));
  }
    
  virtual int n_threads () {return 1;};
  virtual int work_chunk_levels () {return 5;};
  virtual int estimated_levels_needed () {return 0;};
};


vector<T_float> watson_henney_test (int minscat, int maxscat,
				    int n_forced=1, int seed = 0)
{  
  const T_float rho=2;
  const int grid_n=1;
  const T_float emitter_z = -0.5;
  const vec3d grid_extent(1000,1000,0.5);
  const T_float distance = 10000;
  const int n_threads=8;
  const int n_lambda=1;
  const int n_rays=1000000;

  T_lambda lambda(n_lambda);lambda=1;

  mcrx::seed(seed);
  T_unit_map units;
  units ["length"] = "kpc";
  units ["mass"] = "Msun";
  units ["luminosity"] = "W";
  units ["wavelength"] = "m";
  units ["L_lambda"] = "W/m";

  // load grain model
  T_lambda kk(n_lambda); kk=1;
  T_lambda aa(n_lambda); aa=0.5;
  T_lambda gg(n_lambda); gg=0.5;
  T_dust_model::T_scatterer_vector sv;
  sv.push_back(boost::shared_ptr<T_scatterer> 
	       (new simple_HG_dust_grain<polychromatic_scatterer_policy, mcrx_rng_policy> (kk,gg,aa)));
  T_dust_model model (sv);
  model.set_wavelength(lambda);
  
  const T_float albedo = model.get_scatterer(0 ).albedo()(0);

  T_densities rhov(n_lambda);rhov=rho;
  uniform_factory factory (rhov, n_lambda);
  // first create the adaptive grid
  boost::shared_ptr<adaptive_grid<T_grid::T_data> > tempgrid
    (new adaptive_grid<T_grid::T_data>  
     (-grid_extent, grid_extent,
      vec3i (grid_n, grid_n, grid_n), factory));

  // and then the full_sed_grid
  T_grid gr(tempgrid, units);

  cout << "setting up emission" << endl;
  auto_ptr<T_emission> em;
  //em.reset(new pointsource_emission<polychromatic_policy, T_rng_policy> 
  //	   (vec3d (0., 0., emitter_z), lambda));
  em.reset(new laser_emission<polychromatic_policy, T_rng_policy> 
	   (vec3d (0., 0., emitter_z), vec3d(0,0,1), lambda));
  
  vector<pair<T_float, T_float> > campos;
  for(T_float theta=0; theta<=180; theta+=10)
    campos.push_back(make_pair(theta*constants::pi/180,0.));

  auto_ptr<T_emergence> e(new T_emergence(campos, distance, 0.1*distance, 1));

  std::vector<T_rng::T_state> states= generate_random_states (n_threads);
  T_rng_policy rng;

  xfer<T_dust_model, T_grid, T_rng_policy> x (gr, *em, *e, model, rng, 1e-2, 0,n_forced,10., false, false, minscat, maxscat);
  long n_rays_actual = 0;
  dummy_terminator t;
  T_biaser b(0);
  shoot (x, scatter_shooter<T_biaser> (b), t, states, n_rays,
  	 n_rays_actual, n_threads,0, "");

  T_float normalization = 1./n_rays_actual;

  // analyze flux in images
  vector<T_float> result;
  for (T_emergence::iterator i = e->begin(); i != e->end(); ++i) {
    const T_emergence::T_camera::T_image image =
      (*i)->get_image ();
    const T_float scatter = normalization*sum (image)*distance*distance;
    result.push_back(scatter);
    const T_float theta =(*i)->get_theta();
  }
  return result;
}

bool print_test(const vector<vector<T_float> >& res, const array_2& theo,
		int ncol, int roffset, int toffset, bool comp)
{
  T_float maxdiff=0;
  if(comp) cout << "Fractional differences:\n";
  for(int i=0; i<res[0].size();++i) {
    for(int j=0; j<ncol;++j) {
      T_float diff= res[j+roffset][i]!=0 ? 
	res[j+roffset][i]/theo(i,j+toffset)-1 :
	(theo(i,j+toffset)==0 ? 0 : blitz::huge(T_float()));
      if (comp)
	maxdiff=max(maxdiff, abs(diff));

      cout << 
	(comp ? diff : res[j+roffset][i] )
	   << '\t';
    }
    cout << endl;
  }
  cout << endl;
  if(comp) {
    cout << "Max fractional difference: " << maxdiff << endl;
    if(maxdiff>1e-2) {
      cout << "TEST FAILED!\n";
      return false;
    }
    else {
      cout << "Test passed.\n";
      return true;
    }
  }
}

bool test_scattered_intensity ()
{
  ifstream whf("/home/patrik/tests_sunrise/wh01/wh01results.txt");
  vector<T_float> whd((istream_iterator<T_float>(whf)),
			 istream_iterator<T_float>());
  assert(whd.size()==9*19);
  array_2 whdata(&whd[0], shape(19,9), neverDeleteData);

  vector<vector<T_float> > res;
  res.push_back(watson_henney_test (0,0, 0, 42));
  res.push_back(watson_henney_test (1,1, 0, 42));
  res.push_back(watson_henney_test (2,2, 0, 42));
  res.push_back(watson_henney_test (3,blitz::huge(int()), 0, 42));
  res.push_back(watson_henney_test (0,0, 1, 42));
  res.push_back(watson_henney_test (1,1, 1, 42));
  res.push_back(watson_henney_test (2,2, 1, 42));
  res.push_back(watson_henney_test (3,blitz::huge(int()), 1, 42));

  bool pass=true;
  cout << "\n\nNo Forced scattering\n";
  print_test(res,whdata,4,0,5,false);
  pass &=print_test(res,whdata,4,0,5,true);

  cout << "\n\nWith Forced scattering\n";
  print_test(res,whdata,4,4,5,false);
  pass &=print_test(res,whdata,4,4,5,true);

  return pass;
}

int main (int, char**)
{
  return !test_scattered_intensity();
}
